#include "runLen.h"
#include "unionFindSet.h"			// 并查集
#include <omp.h>
#include <cstdint>

// 平台相关的位操作函数
#ifdef _MSC_VER
#include <intrin.h>
#define CTZ _BitScanForward
#else
#define CTZ __builtin_ctz
#endif

#define SIMD 1


/**
* @brief 	使用SIMD并行加速的 基于游程的二值图像连通域标记算法
* @param	_binImg 二值图像
* @param	_lableImg 标记了连通域的二值图像
* @param	offset 连通域偏移量
* @param	regionInf 存放每个区域的游程信息
* @return	连通域数量
*/
int runLen(cv::Mat& _binImg, cv::Mat& _lableImg, int offset,
	std::vector<std::vector<runLenInf>>& regionInf)
{
#if 1
	if (_binImg.empty() || _binImg.type() != CV_8UC1)
	{
		std::cout << "error: _binImg.type() != CV_8UC1 or _binImg is empty" << std::endl;
		return -1;
	}
	_lableImg.release();
	_binImg.convertTo(_lableImg, CV_32SC1);
	cv::copyMakeBorder(_binImg, _binImg, 1, 1, 1, 1, cv::BORDER_CONSTANT, 0);
	cv::copyMakeBorder(_lableImg, _lableImg, 1, 1, 1, 1, cv::BORDER_CONSTANT, 0);
#endif

	// 初始化变量
	int nRuns = 0; // 游程数量

	//----------------------------------------------------------------------------------------------------------------
	// （1）游程获取 开始
	int rows = _binImg.rows;
	int cols = _binImg.cols;

	std::vector<int> stRun;
	std::vector<int> enRun;
	std::vector<int> rowRun;
	std::vector<int> runLabels;

	// 预先分配内存，减少push_back的开销
	size_t runMem = rows * 100;
	stRun.reserve(runMem);		// 预先分配内存给vector，防止反复开辟新内存降低速度
	rowRun.reserve(runMem);		// 只有容器的容量超过了预分配的内存容量runMem，再次push_back才会重新给容器分配内存
	enRun.reserve(runMem);
	runLabels.reserve(runMem);

#if SIMD

	for (int i = 0; i < rows; ++i) {
		const uchar *rowData = _binImg.ptr<uchar>(i);
		int j = 1;

		// SIMD 处理主循环
		for (; j < cols - 16; j += 16) {
			Vec16uc v0(0), v255(255);
			Vec16uc vL = Vec16uc().load(rowData + j - 1); // j-1 到 j+14
			Vec16uc vO = Vec16uc().load(rowData + j);     // j 到 j+15

			Vec16cb vStRun = (vL == v0) & (vO == v255);
			Vec16cb vEnRun = (vL == v255) & (vO == v0);

			// 批量处理 stRun (0->255 跳变)
			uint16_t stMask = to_bits(vStRun);
			while (stMask) {
				unsigned long index;
#ifdef _MSC_VER
				CTZ(&index, stMask);
#else
				index = CTZ(stMask);
#endif
				int pos = j + static_cast<int>(index);
				stRun.push_back(pos);
				rowRun.push_back(i);
				stMask &= stMask - 1;  // 清除最低设置位
			}

			// 批量处理 enRun (255->0 跳变)
			uint16_t enMask = to_bits(vEnRun);
			while (enMask) {
				unsigned long index;
#ifdef _MSC_VER
				CTZ(&index, enMask);
#else
				index = CTZ(enMask);
#endif
				int pos = j + static_cast<int>(index) - 1;
				enRun.push_back(pos);
				enMask &= enMask - 1;  // 清除最低设置位
			}
		}

		// 处理剩余像素（标量方式）
		for (; j < cols; j++) {
			if (rowData[j - 1] == 0 && rowData[j] == 255) {
				stRun.push_back(j);
				rowRun.push_back(i);
			}
			else if (rowData[j - 1] == 255 && rowData[j] == 0) {
				enRun.push_back(j - 1);
			}
		}
	}
#endif

	nRuns = (int)stRun.size(); // 获取游程总数

	// （1）游程获取 结束
	//----------------------------------------------------------------------------------------------------------------


	//----------------------------------------------------------------------------------------------------------------
	// （2）获取等价对 开始
	runLabels.assign(nRuns, 0); // 存储每个游程的标签 初始化标签数组
	std::vector<std::pair<int, int>> equivalences; // 存储等价标签对

	int totalLables = 1;         // 下一个可用的标签号
	int rowCur = 0;              // 当前行号
	int currRowStIdx = 0;        // 当前行的第一个游程索引
	int prevRowStIdx = -1;       // 上一行的第一个游程索引

	for (int i = 0; i < nRuns; i++)
	{
		if (rowRun[i] != rowCur)
		{
			rowCur = rowRun[i];
			prevRowStIdx = currRowStIdx;
			currRowStIdx = i;
		}
		if (prevRowStIdx == -1) // 第一行没有上一行
			continue;
		// 提前计算当前游程的范围，避免重复计算
		int currSt = stRun[i];
		int currEn = enRun[i];
		int currRow = rowRun[i];
		// 对上一行游程进行遍历
		for (int j = prevRowStIdx; j < currRowStIdx; j++)
		{
			// 提前计算上一行游程的范围
			int prevEn = enRun[j];
			int prevSt = stRun[j];
			int prevRow = rowRun[j];
			// 快速检查：行号必须相差1
			if (currRow != prevRow + 1)
				continue;
			// 快速检查：范围是否有重叠
			if (currSt > prevEn + offset || currEn < prevSt - offset)
				continue;
			// 详细检查
			if (currSt <= prevEn + offset && currEn >= prevSt - offset)
			{
				if (runLabels[i] == 0)
					runLabels[i] = runLabels[j];
				else if (runLabels[i] != runLabels[j])
					equivalences.push_back(std::make_pair(runLabels[i], runLabels[j]));
			}
		}
		if (runLabels[i] == 0)
			runLabels[i] = totalLables++;
	}

	totalLables--; // 调整标签总数

	// （2）获取等价对 结束
	//----------------------------------------------------------------------------------------------------------------


	//----------------------------------------------------------------------------------------------------------------
	// （3）等价对求解 开始 - 使用并查集

	// 初始化并查集
	std::vector<int> ufs(totalLables, -1);			// 可优化
	//init_ufs(ufs, totalLables);

	// 使用并查集合并等价对
	for (const auto& eq : equivalences) {
		int label1 = eq.first - 1; // 转换为0-based索引
		int label2 = eq.second - 1;
		Union(ufs.data(), label1, label2);
	}

	// 构建标签映射表
	std::vector<int> labelMap(totalLables, 0);
	int newLabel = 0;

	// 首先找到所有根节点并分配新标签
	int* pUfs = ufs.data();
	int* pLabelMap = labelMap.data();

	for (int i1 = 0; i1 < totalLables; ++i1) {
		if (pUfs[i1] < 0) { // 根节点
			pLabelMap[i1] = ++newLabel;
		}
	}

	// 然后为所有节点分配映射标签
	for (int i2 = 0; i2 < totalLables; ++i2) {
		if (pUfs[i2] >= 0) { // 非根节点
			int root = Find(ufs.data(), i2);
			pLabelMap[i2] = pLabelMap[root];
		}
	}

	// （3）等价对求解 结束
	//----------------------------------------------------------------------------------------------------------------

	//std::cout << "连通域个数：" << newLabel << std::endl;
#if SIMD
	// 更新标签图像 - SIMD优化版本
	for (int i = 0; i < nRuns; ++i)
	{
		int oldLabel = runLabels[i]; // 这是原始标签（1-based）
		int newLabel = labelMap[oldLabel - 1]; // 转换为0-based索引
		// 更新游程的标签
		runLabels[i] = newLabel;

		int start = stRun[i];
		int end = enRun[i];
		int length = end - start + 1;

		int* pData = &_lableImg.ptr<int>(rowRun[i])[start];

		// 使用SIMD向量填充
		Vec16i vLabel(newLabel);  // 创建包含16个label值的向量

		// 每次处理16个像素
		int j = 0;
		for (; j <= length - 16; j += 16)
		{
			vLabel.store(pData + j);  // 一次性存储16个值
		}

		// 处理剩余像素
		for (; j < length; ++j)
		{
			pData[j] = newLabel;
		}
	}
#endif
	//----------------------------------------------------------------------------------------------------------------
	//存储每个区域的全部游程
	regionInf.clear();
	regionInf.resize(newLabel); // 共有：0 - N-1 个连通域，每个连通域标签：1 - N

	// 收集每个连通域的游程信息
	for (int i = 0; i < nRuns; ++i)
	{
		runLenInf runInfo;
		runInfo.stRow = rowRun[i];
		runInfo.stCol = stRun[i];
		runInfo.enCol = enRun[i];
		regionInf[runLabels[i] - 1].push_back(runInfo);
	}
	//----------------------------------------------------------------------------------------------------------------

	//----------------------------------------------------------------------------------------------------------------
	// 内存回收
	{
		std::vector<int>().swap(stRun);
		std::vector<int>().swap(enRun);
		std::vector<int>().swap(rowRun);
		std::vector<int>().swap(runLabels);
		std::vector<std::pair<int, int>>().swap(equivalences);
		std::vector<int>().swap(ufs); // 新增：释放并查集内存
		std::vector<int>().swap(labelMap);
	}
	//----------------------------------------------------------------------------------------------------------------
	return newLabel;
}



